Skip to content

optimize: Optimize rockdb batch query performance#2982

Open
lokidundun wants to merge 16 commits intoapache:masterfrom
lokidundun:improvequery
Open

optimize: Optimize rockdb batch query performance#2982
lokidundun wants to merge 16 commits intoapache:masterfrom
lokidundun:improvequery

Conversation

@lokidundun
Copy link
Copy Markdown
Contributor

@lokidundun lokidundun commented Mar 29, 2026

Purpose of the PR

Main Changes

The core change in this PR is a refactor/optimization of the RocksDB batch query path in three main areas: request coalescing, result backfilling, and shard-level parallelism control.
The focus is not on the API layer, but on how batch get is executed internally and how objects are reused, with the goal of reducing query latency and memory jitter under high concurrency.

Verifying these changes

  • Trivial rework / code cleanup without any test coverage. (No Need)
  • Already covered by existing tests, such as (please modify tests here).
  • Need tests and can be verified as follows:
    • xxx

Does this PR potentially affect the following parts?

Documentation Status

  • Doc - TODO
  • Doc - Done
  • Doc - No Need

@dosubot dosubot bot added size:XS This PR changes 0-9 lines, ignoring generated files. perf store Store module labels Mar 29, 2026
@lokidundun lokidundun closed this Mar 30, 2026
@github-project-automation github-project-automation bot moved this from In progress to Done in HugeGraph PD-Store Tasks Mar 30, 2026
@lokidundun lokidundun reopened this Mar 30, 2026
@github-project-automation github-project-automation bot moved this from Done to In progress in HugeGraph PD-Store Tasks Mar 30, 2026
@lokidundun
Copy link
Copy Markdown
Contributor Author

This CI failure is unrelated to the changes in this PR. The PR focuses on optimizing RocksDB batch query performance, and the failing build check does not involve the code modified here.

@imbajin
Copy link
Copy Markdown
Member

imbajin commented Mar 31, 2026

Already rerun CI (also could check the tests could pass in local env)

@imbajin
Copy link
Copy Markdown
Member

imbajin commented Mar 31, 2026

seems a little strange (after rerun 3 times still failed)

image

@dosubot dosubot bot added size:M This PR changes 30-99 lines, ignoring generated files. and removed size:XS This PR changes 0-9 lines, ignoring generated files. labels Apr 6, 2026
@lokidundun
Copy link
Copy Markdown
Contributor Author

@imbajin could you please take another look when you are convenient❤️

Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR introduces batched backend fetching for queryVerticesByIds() to reduce overhead when querying many vertex ids (e.g., g.V(id1, id2, ...)) by splitting backend id lookups into multiple IdQuery requests.

Changes:

  • Collect backend-only vertex ids during queryVerticesByIds() and issue backend queries in batches using query.batch_size.
  • Add per-batch capacity checks before executing backend queries.
  • Minor formatting adjustment in RocksDB table iterator helper signature.

Reviewed changes

Copilot reviewed 1 out of 2 changed files in this pull request and generated 5 comments.

File Description
hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/tx/GraphTransaction.java Batch backend id queries in queryVerticesByIds() instead of building a single IdQuery for all backend ids.
hugegraph-server/hugegraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBTable.java Formatting-only change to method parameter indentation.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines +809 to +814
if (!backendIds.isEmpty()) {
// Query from backend store
query.mustSortByInput(false);
Iterator<HugeVertex> it = this.queryVerticesFromBackend(query);
QueryResults.fillMap(it, vertices);
final int batch = this.batchSize > 0 ? this.batchSize : backendIds.size();
for (int i = 0; i < backendIds.size(); i += batch) {
int end = Math.min(i + batch, backendIds.size());
IdQuery query = new IdQuery(type);
Copy link

Copilot AI Apr 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The PR title says this is a RocksDB optimization, but the functional change here is backend-agnostic batching in GraphTransaction (affecting all stores, especially RPC-based ones). Please align the PR title (or add RocksDB-specific changes) so the intent matches what’s actually being modified.

Copilot uses AI. Check for mistakes.
Comment on lines 787 to 805
@@ -799,17 +800,30 @@ protected Iterator<Vertex> queryVerticesByIds(Object[] vertexIds, boolean adjace
// Found from local tx
vertices.put(vertex.id(), vertex);
} else {
// Prepare to query from backend store
query.query(id);
// store the IDs queried from backend
backendIds.add(id);
}
Copy link

Copilot AI Apr 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change allocates backendIds as a second full list of ids (in addition to ids), which doubles memory usage for large g.V(id1,id2,...) calls (up to Query.DEFAULT_CAPACITY). Consider batching/issuing backend IdQuery requests incrementally during the main loop (flush when reaching batch size) to avoid retaining all backend ids at once.

Copilot uses AI. Check for mistakes.
Comment on lines +811 to +818
final int batch = this.batchSize > 0 ? this.batchSize : backendIds.size();
for (int i = 0; i < backendIds.size(); i += batch) {
int end = Math.min(i + batch, backendIds.size());
IdQuery query = new IdQuery(type);
for (int j = i; j < end; j++) {
Id id = backendIds.get(j);
query.query(id);
}
Copy link

Copilot AI Apr 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With batching, duplicated ids that fall into different batches will trigger repeated backend reads/RPCs for the same id. You can keep the output behavior (duplicates preserved via ids) while deduplicating backend fetches (e.g., track a seen-set for backendIds or build per-batch unique ids) to avoid redundant backend queries.

Copilot uses AI. Check for mistakes.
query.mustSortByInput(false);
Iterator<HugeVertex> it = this.queryVerticesFromBackend(query);
QueryResults.fillMap(it, vertices);
final int batch = this.batchSize > 0 ? this.batchSize : backendIds.size();
Copy link

Copilot AI Apr 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

final int batch = this.batchSize > 0 ? this.batchSize : backendIds.size(); is effectively always this.batchSize because query.batch_size is configured with a minimum of 1 (see CoreOptions.QUERY_BATCH_SIZE). Removing the dead fallback branch will simplify the code.

Suggested change
final int batch = this.batchSize > 0 ? this.batchSize : backendIds.size();
final int batch = this.batchSize;

Copilot uses AI. Check for mistakes.
Comment on lines +809 to +826
if (!backendIds.isEmpty()) {
// Query from backend store
query.mustSortByInput(false);
Iterator<HugeVertex> it = this.queryVerticesFromBackend(query);
QueryResults.fillMap(it, vertices);
final int batch = this.batchSize > 0 ? this.batchSize : backendIds.size();
for (int i = 0; i < backendIds.size(); i += batch) {
int end = Math.min(i + batch, backendIds.size());
IdQuery query = new IdQuery(type);
for (int j = i; j < end; j++) {
Id id = backendIds.get(j);
query.query(id);
}
// Single batch capacity check
Query.checkForceCapacity(query.idsSize());

// Query from backend store
query.mustSortByInput(false);
Iterator<HugeVertex> it = this.queryVerticesFromBackend(query);
QueryResults.fillMap(it, vertices);
}
Copy link

Copilot AI Apr 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new multi-batch path isn’t covered by tests. Please add a unit/integration test that exercises queryVerticesByIds() with vertexIds.length > query.batch_size, including (1) duplicates across a batch boundary and (2) mixed local-tx + backend ids, to ensure results and NotFoundException behavior remain unchanged.

Copilot uses AI. Check for mistakes.
@lokidundun lokidundun changed the title optimize: Optimize RocksDB batch query performance optimize: Optimize batch query performance Apr 7, 2026
// NOTE: allowed duplicated vertices if query by duplicated ids
List<Id> ids = InsertionOrderUtil.newList();
Map<Id, HugeVertex> vertices = new HashMap<>(vertexIds.length);
Set<Id> fetchedIds = InsertionOrderUtil.newSet();
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ 这里新增 fetchedIds 之后,非相邻的重复 id 也会被全局去重;而旧逻辑里 IdQuery.query() 只会折叠相邻重复 id。最终返回结果看起来应该还是保持重复输出,但真实的后端访问路径已经变了。建议补一个回归测试,至少覆盖 超过 query.batch_size跨 batch 的重复 id缺失 id + checkMustExist 这几个组合场景,避免后面再改这里时把语义悄悄带偏。

Map<Id, HugeVertex> vertices = new HashMap<>(vertexIds.length);
Set<Id> fetchedIds = InsertionOrderUtil.newSet();
IdQuery batchQuery = null;
final int batchSize = this.batchSize;
Copy link
Copy Markdown
Member

@imbajin imbajin Apr 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ 这里把 batching 放在 GraphTransaction 通用层后,会影响所有 backend,而不只是 issue #2674 里提到的 RPC backend。

以 RocksDB 为例,当前 queryByIds() 仍然是逐 id 展开查询,并没有真正走 multi-get;现在强制按 query.batch_size 拆成多个 IdQuery,很可能只是增加额外的 query/iterator 次数。建议把这类分批策略下沉到具体 backend,或者至少通过 feature/store type 把它限定在 HBase/HStore 这类 RPC backend 上,避免把针对性优化变成全局行为变化。

PS: 后续我们应该让 RocksDB 使用上原生的 multi-get API (这应该是之前的 TODO)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. GraphTransaction 不再按 batch_size 拆分 ID,一次性下发完整 IdQuery;
  2. RocksDBTables.Vertex/Edge 已覆写 queryByIds(),在 !session.hasChanges() 时走 multiGetAsList() 原生批量读取,脏 session 安全回退到逐 id scan;

@dosubot dosubot bot removed the size:M This PR changes 30-99 lines, ignoring generated files. label Apr 16, 2026
@dosubot dosubot bot added the size:S This PR changes 10-29 lines, ignoring generated files. label Apr 16, 2026
@lokidundun lokidundun requested a review from imbajin April 17, 2026 08:18
Collection<Id> ids) {
// TODO: use getByIds() after batch version multi-get is ready
if (!session.hasChanges()) {
return this.getByIds(session, ids);
Copy link
Copy Markdown
Member

@imbajin imbajin Apr 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ This RocksDB-specific optimization looks reasonable to me as a focused improvement on its own. One thing that may be worth clarifying is how HStore should be handled, since it is the main distributed backend we maintain and its batch-query path is a separate concern from RocksDB.

If you think it makes sense, we could either cover HStore in this PR as well, or keep this PR scoped to RocksDB and follow up with a separate PR for HStore so the behavior and performance trade-offs can be reviewed independently.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! Let's keep this PR focused on RocksDB for now, and I'll handle HStore in a separate follow-up PR.

protected BackendColumnIterator queryByIds(RocksDBSessions.Session session,
Collection<Id> ids) {
if (!session.hasChanges()) {
return this.getByIds(session, ids);
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Since this adds a new multi-get path for vertex/edge id queries, it would be great to add a small RocksDB regression test here. The core cases that seem worth covering are:

  1. batch query with multiple existing ids
  2. batch query with existing + missing ids mixed together
  3. duplicate ids in the input
  4. fallback to the old path when session.hasChanges() is true

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Testing the fallback path (session.hasChanges() = true) is not feasible in unit tests, as RocksDBStdSessions asserts !this.hasChanges() on all read operations (get/scan) when assertions are enabled.

@dosubot dosubot bot added size:L This PR changes 100-499 lines, ignoring generated files. and removed size:S This PR changes 10-29 lines, ignoring generated files. labels Apr 17, 2026
@lokidundun lokidundun changed the title optimize: Optimize batch query performance optimize: Optimize rockdb batch query performance Apr 18, 2026

protected BackendColumnIterator getByIds(RocksDBSessions.Session session, Set<Id> ids) {
protected BackendColumnIterator getByIds(RocksDBSessions.Session session,
Collection<Id> ids) {
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

‼️ Set<Id>Collection<Id> 语义变化需注意

原方法接收 Set<Id>(天然去重),改为 Collection<Id> 后,传入 List 时若含重复 ID,RocksDB multiGet 会对同一 key 重复查询并返回重复结果。

测试 testVertexQueryByIdsWithDuplicateIds 验证了这个行为(id1 返回 2 次),但这与原 Set 语义不一致。需要确认上层 IdQuery 的 ids 是否可能含重复——如果含重复,行为变更可能导致上层重复处理数据。

建议:在 getByIds 入口去重以保持原语义,或者明确文档说明 Collection 含重复返回的行为变更。

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added LinkedHashSet deduplication at the getByIds() entry point while preserving input order; skips if Set is already passed. Test updated to verify dedup behavior. The
backend deduplication does not affect final semantics since GraphTransaction reassembles results by the original input order.

* which have a pre-existing assertion `assert !this.hasChanges()` in
* RocksDBStdSessions. This assertion is disabled in production but fires
* during unit tests when assertions are enabled. The dispatch logic itself
* is covered by the implementation in RocksDBTables.Vertex/Edge.queryByIds().
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

‼️ fallback 路径(hasChanges() == true)缺乏测试覆盖

当 session 有未提交变更时走 super.queryByIds()(scan 路径),但测试中因 RocksDBStdSessionsassert !this.hasChanges() 断言无法覆盖此场景。

核心分支逻辑(if/else)只测了一半。建议考虑用 mock session 或在集成测试中补充覆盖 hasChanges() == true 的场景。

}
return super.queryByIds(session, ids);
}
}
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Vertex 和 Edge 的 queryByIds 逻辑完全重复

Vertex.queryByIds()Edge.queryByIds() 代码完全一致(5 行相同逻辑)。可以考虑:

  1. 提取到父类 RocksDBTable.queryByIds() 中统一做 hasChanges() 判断
  2. 或者提供一个公共的 helper 方法

这样可以减少重复,也降低后续维护时两处不一致的风险。

Collection<Id> ids) {
// TODO: use getByIds() after batch version multi-get is ready
if (!session.hasChanges()) {
return this.getByIds(session, ids);
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ fallback 路径的实际效果需确认

hasChanges() == true 时,fallback 到 super.queryByIds() 会对每个 ID 调用 this.queryById()。而 Vertex/Edge 已经 override queryByIdgetById(point get),所以 fallback 路径实际也会调用 getById

RocksDBStdSessionsget() 方法有 assert !this.hasChanges() 断言——也就是说 fallback 路径在开启断言的环境下依然会触发 assert 失败。

需要确认:hasChanges() == true 时是否真的能安全走 fallback?还是说这个分支在 RocksDB 后端实际上永远不会被触发?

@@ -182,7 +182,9 @@
@Override
protected BackendColumnIterator queryByIds(RocksDBSessions.Session session,
Collection<Id> ids) {
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 PR 描述与实际改动不完全匹配

PR 描述提到 "improves the performance of Gremlin queries... when using RPC-based backends such as HBase and HStore",但代码改动全在 hugegraph-rocksdb 模块。建议更新 PR 描述,准确反映本次优化的范围是 RocksDB 后端。

import org.apache.hugegraph.unit.id.SplicingIdGeneratorTest;
import org.apache.hugegraph.unit.mysql.MysqlUtilTest;
import org.apache.hugegraph.unit.mysql.WhereBuilderTest;
import org.apache.hugegraph.unit.rocksdb.RocksDBTableQueryByIdsTest;
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 import 未按字母序排列

RocksDBTableQueryByIdsTest 的 import 插在 RocksDBCountersTest 之前,不符合字母顺序。建议调整为:

Suggested change
import org.apache.hugegraph.unit.rocksdb.RocksDBTableQueryByIdsTest;
import org.apache.hugegraph.unit.rocksdb.RocksDBCountersTest;
import org.apache.hugegraph.unit.rocksdb.RocksDBSessionTest;
import org.apache.hugegraph.unit.rocksdb.RocksDBSessionsTest;
import org.apache.hugegraph.unit.rocksdb.RocksDBTableQueryByIdsTest;

/**
* Subclass that exposes the protected queryByIds for testing.
*/
private static class TestEdgeTable extends RocksDBTables.Edge {
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 只测了 out edge 方向

TestEdgeTable 固定传 true(out edge),in edge 方向未覆盖。虽然逻辑相同,但建议至少补一个 in 方向的 case 或在注释中说明为什么只测 out。

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

perf size:L This PR changes 100-499 lines, ignoring generated files. store Store module

Projects

Status: In progress

Development

Successfully merging this pull request may close these issues.

[Improve] Gremlin query ids batch query to backend such as HBase/HStore

4 participants